#!/usr/bin/python
'''
Created by Colin Gillespie <c.gillespie@ncl.ac.uk>
LastChangedDate: 2008-03-12 19:52:33 +0000 (Wed, 12 Mar 2008) $

Constructs the reaction information

'''
from __future__ import division
__revision__ = '$Rev: 125 $'


from copy import copy
import os
import libsbml
from sympy import *
 

class ReactionInformation(object):
    '''
        Store info about the reaction
    '''
    
    def __init__(self, sbml_id=None, species_order=[]):
        '''Set reaction id'''
        self.sbml_id = sbml_id
        self.stoic = None
        self.species_order = species_order
       
        #alpha==numerator, beta= denominator.
        self.alpha = None
        self.beta = None
        
        #a_left refers to P(x-s) and a_right to P(x)
        self.a_left = None
        self.a_right = None
        
        
    def setStoic(self, stoic):
        '''Store stoichiometry of the reaction'''
        self.stoic = stoic
        return True
    
    def initialise(self):
        del self.species_order[:]
        
    def setSpeciesOrder(self, new_order):
        self.initialise()
        
        for sp in new_order:
            self.species_order.append(sp)
        return True
    
    def setAlpha(self, alpha_dict):
        '''The numerator of the rate law'''
        self.alpha = alpha_dict   
   
        
    def setBeta(self, beta_dict):
        '''The numerator of the rate law'''
        self.beta = beta_dict   


class GetStoichiometryDict(object):
        
    def __init__(self, sbml_string=None, sbml_file=None, dir_in=None):
        #Handle the SBML
        r = libsbml.SBMLReader()

        if sbml_string:
            self.d = r.readSBMLFromString(sbml_string)
        else:
            if dir_in == None:
                dir_in = os.getcwd()
            file_path = os.path.join(dir_in, sbml_file)
            assert os.path.exists(file_path), "Invalid path"
            self.d = r.readSBML(file_path)
        
        self.d.setLevelAndVersion(2, 4)
        self.m = self.d.getModel()
        self.pa_dict = {}
        self.index2species = {}
        self.hasDenominator = False
       
        #Create species dictionaires
        self.sp_dict, self.empty_list = self.getSpeciesDict()
        
        #Create parameter dictionaires

        self.updateParameterDict(self.m.getListOfParameters())
        self.reaction_list = []
        self.sps_as_symbols = None
        
        self.addReactions()
        self.beta_prod = None
        self.beta_prod_plus = None
        self.beta_prod_minus = None

        self.MakeBetaProducts()
        self.h = (self.beta_prod * self.beta_prod_minus).as_dict()
        self.makeFVector()
        #self.f = None
        #self.g = None
        
         
        
    def MakeBetaProducts(self):
        
        sym_sps = self.sps_as_symbols

        sp_stoic_list_plus = []
        sp_stoic_list_minus = []
        
        for i, reaction in enumerate(self.reaction_list):
            sp_stoic_entry_minus = {}
            sp_stoic_entry_plus = {}
            print 'STOIC', reaction.stoic
            for sp, stoic in zip(sym_sps, reaction.stoic):
                if stoic != 0:
                    sp_stoic_entry_minus[sp] = sp-stoic
                    sp_stoic_entry_plus[sp] = sp+stoic
            sp_stoic_list_plus.append(sp_stoic_entry_plus)
            sp_stoic_list_minus.append(sp_stoic_entry_minus)
        
        self.beta_prod = 1
        self.beta_prod_minus = 1
        self.beta_prod_plus = 1        
        for i, reaction in enumerate(self.reaction_list):
            self.beta_prod *= reaction.beta
            if reaction.beta != reaction.beta.subs(sp_stoic_list_minus[i]):
                print 'hi'
                self.beta_prod_minus *= reaction.beta.subs(sp_stoic_list_minus[i])
            self.beta_prod_plus *= reaction.beta.subs(sp_stoic_list_plus[i])
            print 'XXXXXXXXXXXXXXXX'
            print self.beta_prod, self.beta_prod_minus, self.beta_prod_plus
            print 'XXXXXXXXXXXXXXXX'
        
        for i, reaction in enumerate(self.reaction_list):
            
            a_left = self.beta_prod*reaction.alpha.subs(sp_stoic_list_minus[i])
            print 0, a_left
            a_left = a_left.div_term(reaction.beta.subs(sp_stoic_list_minus[i])*self.beta_prod_minus)
            print 1, a_left
            print 2, a_left.map_coeffs(simplify) 
            
            #a_left = reaction.alpha.subs(sp_stoic_list_minus[i])*self.beta_prod
            
            #a_left *= self.beta_prod_plus/reaction.beta
            # a_left = a_left.subs(sp_stoic_list_plus[i])
            
            a_right = reaction.alpha*self.beta_prod * self.beta_prod_minus/reaction.beta
            reaction.a_left = a_left.as_dict()
            reaction.a_right = a_right.as_dict()
            
        return True

    def makeFVector(self):
        '''This should be move into MakeBetaProducts'''
        return


    def getSpeciesDict(self):
        '''
        returns a dictionary of species, of the form:
        sp_name:{'amount':amount, 'index':index}
        empty_list is length(sp_name) vector of zeros
        '''
        sp_dict = {}
        empty_list = []
        no_of_constant_sp = 0 
        for i, sp in enumerate(self.m.getListOfSpecies()):
            if sp.getConstant():
                self.__addParameter(sp.getId(), int(sp.getInitialAmount()))
                no_of_constant_sp += 1
            else:               
                sp_dict[sp.getId()] = {'amount':int(sp.getInitialAmount())}
                sp_dict[sp.getId()]['index'] = i - no_of_constant_sp
                sp_dict[sp.getId()]['name'] = sp.getName()
                self.index2species[i - no_of_constant_sp] = sp.getId()
                empty_list.append(0)
        return sp_dict, empty_list

    def updateParameterDict(self, pa_list):
        '''
        Scans the global parameters and returns a dictionary
        '''
        for i in pa_list:
            if self.pa_dict.has_key(i.getId()):
                raise 'All parameters must be unique'
            else:
                self.__addParameter(i.getId(), i.getValue())
        return True
    
    def __addParameter(self, pa_id, pa_value):
        self.pa_dict[pa_id] = {'value':pa_value, 'symbol': Symbol(pa_id)}
        return True
        
    def getReactionList(self):
        return self.reaction_list
    
    def addReactions(self):
        '''
            Used to update the delta dict
        '''
        #Get the order of the species
        species_order = copy(self.empty_list)
        for sp, values in self.sp_dict.items():
            species_order[values['index']] = sp

        r = ReactionInformation()
        r.setSpeciesOrder(species_order)

        self.sps_as_symbols = [Symbol(sp) for sp in species_order]
        
        for reaction in self.m.getListOfReactions():
            #Initialise values
            self.alpha_dict = {}
            stoic_values = copy(self.empty_list)
            
            for reactant in reaction.getListOfReactants():
                sp = reactant.getSpecies()
                if self.sp_dict.has_key(sp):
                    r_st = -int(reactant.getStoichiometry())
                    stoic_values[self.sp_dict[sp]['index']] = r_st
                
            for product in reaction.getListOfProducts():
                sp = product.getSpecies()
                if self.sp_dict.has_key(sp):
                    p_st = int(product.getStoichiometry())
                    stoic_values[self.sp_dict[sp]['index']] += p_st
          
            
            #Convert the kinetic law in swiginac
            kinetic_law = reaction.getKineticLaw()
            self.updateParameterDict(kinetic_law.getListOfParameters())
            kl = sympify(kinetic_law.getFormula())
            
            
            #Create a ReactionInformation class
            r = ReactionInformation(reaction.getId())
            r.setStoic(stoic_values)

            num, denom = kl.as_numer_denom()
            #Split the rate law in coeff's of species
            p = Poly(num, *self.sps_as_symbols)
            r.setAlpha(p)
            self.alpha_dict = p.as_dict() 

            
            p = Poly(denom, *self.sps_as_symbols)
            r.setBeta(p)
            self.beta_dict = p.as_dict()   
            if not self.hasDenominator and p!=1:
                self.hasDenominator = True



            self.reaction_list.append(r)
 

 
if __name__ == '__main__':
    #r = ReactionInformation('test')
    #r.setSpeciesOrder(['X','Y'])


    # r.species_order
    #r1 = ReactionInformation('test1')
    #print r1.species_order
    
#    gdd = GetStoichiometryDict(sbml_file = 'xml_files/test/imdeath.xml')
    gdd = GetStoichiometryDict(sbml_file = 'xml_files/test/mm.xml')    
   # gdd = GetStoichiometryDict(sbml_file = 'xml_files/test/rational_bi.xml')
    reaction_list = gdd.getReactionList()
    for value in reaction_list:
        #print 1, value.sbml_id, value.species_order, value.stoic
        
        print 1, value.alpha
        print 2, value.beta
        print 3, value.a_left, value.a_right

#    print gdd.denominator#
#    x=gdd.getReactionList()[0]
#    print x, dir(x), x.alpha, x.getAlpha()







